#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/8/5 3:09 下午
# @Author  : wancheng.b
# @File    : locust_http.py
# @Software: PyCharm
import os
import re

from locust import HttpUser, task, TaskSet, between


class TaskTest(TaskSet):


    def on_start(self):
        '''

        :return:
        '''
        print('TaskTest on_start')
        # 登录的部分可以写到on_start()里面，HttpUser会保存session状态，后面直接调用self.client就会有登录状态
        data = {
            'userName': 'bianwancheng',
            'credentials': 'abc123456d',
            'referrerUrl': 'http://test-admin.51tonglu.com/login'
        }
        headers = {
            'Content-Type': 'application/x-www-form-urlencoded'
        }
        # 可以理解成路由
        response = self.client.post(url='/login.json', data=data, headers=headers)
        print(response.text)
        ticket = re.findall('=(.+?)%', response.json()['data']['location'])
        token_res = self.client.get(url=f'http://test-admin.51tonglu.com/api/sso/token?ticket={ticket[0]}')
        globals()['token'] = token_res.json()['result']['token']


    @task(1)
    def home(self):
        headers = {'token': globals()['token']}
        response = self.client.get(
            'http://test-admin.51tonglu.com/api/companies?page=1&pageSize=10&keyword=%E9%9B%86%E7%BB%9F%E6%9C%88%E4%BB%98%E8%BE%B9',
            headers=headers)

        print('---------------')
        print(response.json())

    # @task(1)
    # def stop_test(self):
    #     self.interrupt()

class MyUser(HttpUser):
    wait_time = between(1, 2)
    tasks = {TaskTest: 1}

    def on_start(self):
        print('httpuser on_start')


if __name__ == '__main__':
    os.system('locust -f locust_http.py --host=http://test.upm.shebao.net')
